import torchvision
import torchvision.transforms as transforms
import os
from tqdm import tqdm
import pandas as pd
import numpy as np
from PIL import Image
# 加载指标计算方法
from utils import *
import pyiqa
from FFT import *
if torch.cuda.is_available():
    device = 'cuda'
else:
    device = 'cpu'
print('Using device:', device)

# 创建pyiqa指标对象
brisque_metric = pyiqa.create_metric('brisque',device=device)
ilniqe_metric = pyiqa.create_metric('ilniqe',device=device)

# 设置数据集的路径和下载选项
Dataset_dir = '../data'
Output_dir = './results_filtered'
train_metrics_dir = os.path.join(Output_dir, 'train_metrics_fitered.csv')
test_metrics_dir = os.path.join(Output_dir, 'test_metrics_filtered.csv')
#如果输出目录不存在，则创建
if not os.path.exists(Output_dir):
    os.makedirs(Output_dir)

#定义数据预处理
transform = transforms.Compose([transforms.ToTensor()]) 
# 读取CITAR-100数据集
trainset = torchvision.datasets.CIFAR100(root=Dataset_dir, train=True, download=True, transform=transform)
testset = torchvision.datasets.CIFAR100(root=Dataset_dir, train=False, download=True, transform=transform)

def generate_img_filtered_metrics(cifar100_dataset, Output_dir):
    print("Calculating image metrics...")
    # 定义将张量转换为PIL图像的函数
    to_pil = transforms.ToPILImage()
    results_list = []  # 使用列表来收集数据

    # 使用tqdm来显示进度
    for i in tqdm(range(len(cifar100_dataset))):
        
        # 读取图像数据(张量形式)
        image_ten, label = cifar100_dataset[i]
        # 将张量转为PIL图像
        image_img = to_pil(image_ten)
        gray_image_img = image_img.convert('L')  # 将PIL图像转为灰度图像
        gray_img_low, gray_img_high = filter_image(gray_image_img, filter_size=2, show_filters=False)
        # # 将PIL图像转为灰度图像
        # gray_img_low = low_img.convert('L') 
        # gray_img_high = high_img.convert('L') 

        # 将灰度图像转为numpy数组
        img_int_low = np.array(gray_img_low).astype(np.int32)
        img_double_low = np.array(gray_img_low).astype(np.float32)
        img_int_high = np.array(gray_img_high).astype(np.int32)
        img_double_high = np.array(gray_img_high).astype(np.float32)

        image_tensor=image_ten.unsqueeze(0)
        # 计算pyiqa指标
        brisque_score = brisque_metric(image_tensor.to(device)).item()
        #ilniqe_score = ilniqe_metric(image_tensor.to(device)).item()


        # 计算信息熵EN
        EN_low = calculate_EN(img_int_low)
        EN_high = calculate_EN(img_int_high)
        # 计算空间频率SF
        SF_low = calculate_SF(img_double_low)
        SF_high = calculate_SF(img_double_high)
        # 计算标准差SD
        SD_low = calculate_SD(img_double_low)
        SD_high = calculate_SD(img_double_high)
        # 计算平均梯度
        AG_low = calculate_AG(img_double_low)
        AG_high = calculate_AG(img_double_high)
        # 计算信噪比SNR
        SNR_low = calculate_SNR(image_ten.to(device))
        SNR_high = calculate_SNR(image_ten.to(device))
        # 将结果添加到列表中
        results_list.append({
            'idx':i,
            'cifar100_label': label,
            'SNR_low': SNR_low,
            'SNR_high': SNR_high,
            'EN_low': EN_low,
            'EN_high': EN_high,
            'SF_low': SF_low,
            'SF_high': SF_high,
            'SD_low': SD_low,
            'SD_high': SD_high,
            'AG_low': AG_low,
            'AG_high': AG_high,
            'BRISQUE': brisque_score,
            # 'BRISQUE_low': brisque_score_low,
            # 'BRISQUE_high': brisque_score_high,
            #'ILNIQE': ilniqe_score,
        })

    print("Saving file...")
    # 将列表转换为DataFrame，并指定列名

    results_df = pd.DataFrame(results_list)

    # 保存结果到csv文件中   
    csv_file_path = Output_dir
    results_df.to_csv(csv_file_path,index=False)
    print("Done.")

if __name__ == '__main__':
    generate_img_filtered_metrics(trainset,train_metrics_dir)
    generate_img_filtered_metrics(testset,test_metrics_dir)